from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import datetime
import MySQLDataFeed as mdf
import SimpleMySqlClass as smc

import backtrader as bt


class TALibStrategy(bt.Strategy):
    params = (
        ("ind", "sma"),
        ("doji", True),
    )

    INDS = [
        "sma",
        "ema",
        "stoc",
        "rsi",
        "macd",
        "bollinger",
        "aroon",
        "ultimate",
        "trix",
        "kama",
        "adxr",
        "dema",
        "ppo",
        "tema",
        "roc",
        "williamsr",
    ]

    def __init__(self):
        if self.p.doji:
            bt.talib.CDLDOJI(
                self.data.open, self.data.high, self.data.low, self.data.close
            )

        if self.p.ind == "sma":
            bt.talib.SMA(self.data.close, timeperiod=25, plotname="TA_SMA")
            bt.indicators.SMA(self.data, period=25)
        elif self.p.ind == "ema":
            bt.talib.EMA(timeperiod=25, plotname="TA_SMA")
            bt.indicators.EMA(period=25)
        elif self.p.ind == "stoc":
            bt.talib.STOCH(
                self.data.high,
                self.data.low,
                self.data.close,
                fastk_period=14,
                slowk_period=3,
                slowd_period=3,
                plotname="TA_STOCH",
            )

            bt.indicators.Stochastic(self.data)

        elif self.p.ind == "macd":
            bt.talib.MACD(self.data, plotname="TA_MACD")
            bt.indicators.MACD(self.data)
            bt.indicators.MACDHisto(self.data)
        elif self.p.ind == "bollinger":
            bt.talib.BBANDS(self.data, timeperiod=25, plotname="TA_BBANDS")
            bt.indicators.BollingerBands(self.data, period=25)

        elif self.p.ind == "rsi":
            bt.talib.RSI(self.data, plotname="TA_RSI")
            bt.indicators.RSI(self.data)

        elif self.p.ind == "aroon":
            bt.talib.AROON(self.data.high, self.data.low, plotname="TA_AROON")
            bt.indicators.AroonIndicator(self.data)

        elif self.p.ind == "ultimate":
            bt.talib.ULTOSC(
                self.data.high, self.data.low, self.data.close, plotname="TA_ULTOSC"
            )
            bt.indicators.UltimateOscillator(self.data)

        elif self.p.ind == "trix":
            bt.talib.TRIX(self.data, timeperiod=25, plotname="TA_TRIX")
            bt.indicators.Trix(self.data, period=25)

        elif self.p.ind == "adxr":
            bt.talib.ADXR(
                self.data.high, self.data.low, self.data.close, plotname="TA_ADXR"
            )
            bt.indicators.ADXR(self.data)

        elif self.p.ind == "kama":
            bt.talib.KAMA(self.data, timeperiod=25, plotname="TA_KAMA")
            bt.indicators.KAMA(self.data, period=25)

        elif self.p.ind == "dema":
            bt.talib.DEMA(self.data, timeperiod=25, plotname="TA_DEMA")
            bt.indicators.DEMA(self.data, period=25)

        elif self.p.ind == "ppo":
            bt.talib.PPO(self.data, plotname="TA_PPO")
            bt.indicators.PPO(self.data, _movav=bt.indicators.SMA)

        elif self.p.ind == "tema":
            bt.talib.TEMA(self.data, timeperiod=25, plotname="TA_TEMA")
            bt.indicators.TEMA(self.data, period=25)

        elif self.p.ind == "roc":
            bt.talib.ROC(self.data, timeperiod=12, plotname="TA_ROC")
            bt.talib.ROCP(self.data, timeperiod=12, plotname="TA_ROCP")
            bt.talib.ROCR(self.data, timeperiod=12, plotname="TA_ROCR")
            bt.talib.ROCR100(self.data, timeperiod=12, plotname="TA_ROCR100")
            bt.indicators.ROC(self.data, period=12)
            bt.indicators.Momentum(self.data, period=12)
            bt.indicators.MomentumOscillator(self.data, period=12)

        elif self.p.ind == "williamsr":
            bt.talib.WILLR(
                self.data.high, self.data.low, self.data.close, plotname="TA_WILLR"
            )
            bt.indicators.WilliamsR(self.data)


def runstrat(args=None):
    args = parse_args(args)

    cerebro = bt.Cerebro()

    cerebro.addstrategy(TALibStrategy, ind=args.ind, doji=not args.no_doji)

   
    sql = f"select sd_date, sd_open, sd_high, sd_low, sd_close, sd_volume from t_stock_k_data"
    data = mdf.MySQLDataFeed(sql)

    # Add the Data Feed to Cerebro
    data0 = cerebro.adddata(data, name="talib")

    cerebro.adddata(data0)

    cerebro.run(runcone=not args.use_next, stdstats=False)
    if args.plot:
        pkwargs = dict(style="candle")
        if args.plot is not True:  # evals to True but is not True
            npkwargs = eval("dict(" + args.plot + ")")  # args were passed
            pkwargs.update(npkwargs)

        cerebro.plot(**pkwargs)


def parse_args(pargs=None):

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Sample for sizer",
    )

    parser.add_argument(
        "--ind",
        required=False,
        action="store",
        default=TALibStrategy.INDS[0],
        choices=TALibStrategy.INDS,
        help=("Which indicator pair to show together"),
    )

    parser.add_argument(
        "--no-doji",
        required=False,
        action="store_true",
        help=("Remove Doji CandleStick pattern checker"),
    )

    parser.add_argument(
        "--use-next",
        required=False,
        action="store_true",
        help=("Use next (step by step) " "instead of once (batch)"),
    )

    # Plot options
    parser.add_argument(
        "--plot",
        "-p",
        nargs="?",
        required=False,
        metavar="kwargs",
        const=True,
        help=(
            "Plot the read data applying any kwargs passed\n"
            "\n"
            "For example (escape the quotes if needed):\n"
            "\n"
            '  --plot style="candle" (to plot candles)\n'
        ),
    )

    if pargs is not None:
        return parser.parse_args(pargs)

    return parser.parse_args()


if __name__ == "__main__":
    runstrat()
